package com.peng.leetcode.backtracking;

import java.util.*;

/**
 * SolveNQueens
 *
 * 51. N 皇后
 *
 * @author: lupeng6
 * @create: 2021/2/4 12:11
 */
public class SolveNQueens {

    public static void main(String[] args) {
        List<List<String>> lists = new SolveNQueens().solveNQueens(3);
        for (int i = 0; i < lists.size(); i++) {
            System.out.println("---------------------");
            for (String s : lists.get(i)) {
                System.out.println(s);
            }
        }
    }

    private Set<Integer> delCols = new HashSet<>();
    private Set<Integer> delDia1 = new HashSet<>();
    private Set<Integer> delDia2 = new HashSet<>();

    private List<int[]> answer = new ArrayList<>();

    private int n;

    public List<List<String>> solveNQueens(int n) {
        this.n = n;
        int[] index = new int[n];
        Arrays.fill(index, -1);
        dfs(0, index);
        return genRes();
    }

    private void dfs(int rowIndex, int[] index) {
        if (rowIndex == n) {
            int[] copy = Arrays.copyOf(index, index.length);
            answer.add(copy);
            return;
        }

        for (int i = 0; i < n; i++) {
            int dia1 = rowIndex + i;
            int dia2 = rowIndex - i;
            if (delCols.contains(i)) {
                continue;
            }
            if (delDia1.contains(dia1)) {
                continue;
            }
            if (delDia2.contains(dia2)) {
                continue;
            }

            delCols.add(i);
            delDia1.add(dia1);
            delDia2.add(dia2);
            index[rowIndex] = i;
            dfs(rowIndex + 1, index);

            // 回溯
            delCols.remove(i);
            delDia1.remove(dia1);
            delDia2.remove(dia2);
            index[rowIndex] = -1;
        }

    }

    private List<List<String>> genRes() {
        int size = answer.size();
        List<List<String>> res = new ArrayList<>(size);
        for (int i = 0; i < size; i++) {
            int[] ints = answer.get(i);
            List<String> item = new ArrayList<>();
            for (int j = 0; j < ints.length; j++) {
                StringBuilder builder = new StringBuilder();
                int index = ints[j];
                for (int k = 0; k < n; k++) {
                    if (k == index) {
                        builder.append("Q");
                    } else {
                        builder.append(".");
                    }
                }
                item.add(builder.toString());
            }

            res.add(item);
        }
        return res;
    }
}
